/*             procpy - Python's extenstion to procps library
 *                             
 * Author: Wojciech Walczak <wojtek.gminick.walczak@gmail.com>
 * Copyright 2008 Wojciech Walczak
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "_procpy.h"

static PyObject *ProcpyError;

/* openproc() flags */
int setflags(void)
{
   int flags;
   
   flags =  PROC_FILLUSR;
   flags |= PROC_FILLGRP;
   flags |= PROC_FILLCOM;
   flags |= PROC_FILLENV;
   flags |= PROC_FILLSTAT;
   flags |= PROC_FILLSTATUS;
   flags |= PROC_FILLMEM;
   flags |= PROC_FILLARG;
   flags |= PROC_FILLWCHAN;

   return flags;
}

/* returns command's name when its argv is not available */
char *
getstatus(int pid)
{
   char fajl[64]="", buf[256]="", *name, *err = "[?]";
   FILE *fp;

   if((/*(char*)*/name = (char*)malloc(256))==NULL)
      return err;

   snprintf(fajl, sizeof(fajl)-1, "/proc/%d/status", pid);
   fajl[sizeof(fajl)-1]='\0';

   if((fp = fopen(fajl,"r"))==NULL)
      return err;

   if(fgets(buf, sizeof(buf)-1, fp)==NULL)
      return err;

   sscanf(buf, "Name:\t%80s", name);
   return name;
}

/* fillcmdline returns argv of the process as a tuple */
PyObject *
fillcmdline(char **cmd, int tid)
{
    int i;
    PyObject *t, *s;

    if((t = PyTuple_New(1))==NULL)
       return NULL;

    if(cmd) {
       for(i=0; cmd[i]; i++) {
          if(_PyTuple_Resize(&t, i+1) != 0)
             return NULL;
          if(!(s = PyString_FromString(cmd[i])))
             return NULL;
          PyTuple_SetItem(t, i, s);
       }
    }
    else {
       s = PyString_FromFormat("[%s]", getstatus(tid));
       PyTuple_SetItem(t, 0, s);
    }

    return t;
}

/* returns a dictionary with environ pairs (i.e. HOME=/home/gminick)
 * as dictionary's key and value (ie. {'HOME':'/home/gminick', ...}) */
PyObject *
getenviron(proc_t *buf)
{
   PyObject *dict = PyDict_New();
   PyObject *key, *value;
   char **eptr = buf->environ;

   if(buf->environ==NULL)
       return dict;

   /* KDE way of setting environment seems a bit weird. 
      Here's the strace output for previous version of procpy:

    [...]
    open("/proc/3067/environ", O_RDONLY)    = 26
    read(26, "\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"..., 2047) = 666
    close(26)                               = 0
    --- SIGSEGV (Segmentation fault) @ 0 (0) ---
    +++ killed by SIGSEGV +++
    
    $ ps auxw|grep 3067
    gminick   3067  0.0  0.6  25048  7052 ?        Ss   21:55   0:00 kdeinit Running...
    $

    So, when /proc/PID/environ starts with '\0', just return empty dictionary.
   */
   if(buf->environ[0][0]=='\0')
       return dict;

   for(; *eptr; eptr++) {
      int i, kvset=0;
      int eptrlen = strlen(*eptr);
      char *tmp = malloc(eptrlen+1);

      if(tmp == NULL)
         return NULL;

      strncpy(tmp, *eptr, eptrlen);
      tmp[eptrlen]='\0';

      for(i=0;i<eptrlen;i++) {
         if(tmp[i] == '=') {
             if((value = PyString_FromString(tmp+i+1)) == NULL)
                return NULL;
             tmp[i]='\0';
             if((key = PyString_FromString(tmp)) == NULL)
                return NULL;
             kvset=1;
         }
      }
      
      /* we have found something in /proc/PID/environ, but the syntax was bad
       * (i.e. there was no '=' sign) so we are creating a key 'rubbish'
       * and assiging the contents of /proc/PID/environ to it.
       *
       * A real-life example:
       * I'm running KDE's konsole in my GNOME environment.
       * konsole starts some KDE processes including dcopserver.
       * The thing is that dcopserver's environ is overwritten
       * by dcopserver's argv. I can repeat that by running kdeinit.
       * Anyway, in this case - we just have to deal with rubbish!
       *
       * $ cat /proc/3071/environ
       * d --suicide
       * $ ps aux|grep 3071
       * gminick   3071  [...]  dcopserver [kdeinit] --nosid --suicide
       * $ 
       */
      if(!kvset && strlen(tmp)>0) {
         if((key = PyString_FromString("rubbish")) == NULL)
            return NULL;
         if((value = PyString_FromString(tmp)) == NULL)
            return NULL;
         kvset=1;  
      }

      if(kvset) {
         if(PyDict_SetItem(dict, key, value) < 0)
            return NULL;
         Py_DECREF(key);
         Py_DECREF(value);
      }
      free(tmp); 
   }
   return dict;
}


/* based on procps's ps (pr_pcpu() in output.c);
   sets buf->pcpu AND returns string object with value in ##.# format */
PyObject *
setpcpu(proc_t *buf)
{
   unsigned long seconds_since_boot = uptime(0,0);
   unsigned long long used_jiffies;
   unsigned long long avail_jiffies;
   PyObject *pcpu;
   
   buf->pcpu = 0;

   used_jiffies = buf->utime + buf->stime;
   avail_jiffies = seconds_since_boot - buf->start_time / Hertz;

   if(avail_jiffies)
      buf->pcpu = (used_jiffies * 1000ULL / Hertz) / avail_jiffies;

   if(buf->pcpu > 999U)
      pcpu = PyString_FromFormat("%u", buf->pcpu/10U);
   else
      pcpu = PyString_FromFormat("%u.%u", buf->pcpu/10U, buf->pcpu%10);

   return pcpu;
}

/* based on pr_pmem() from output.c (procps's ps)
   returns %PMEM as string object in ##.# format */
PyObject*
getpmem(proc_t *buf)
{
   unsigned long pmem = 0;

   meminfo(); // fills kb_main_total;

   pmem = buf->vm_rss * 1000ULL / kb_main_total;

   if(pmem > 999)
      pmem = 999;

   return PyString_FromFormat("%2u.%u", (unsigned)(pmem/10),
                                        (unsigned)(pmem%10));
}

/* based on pr_bsdstart (output.c from procps's ps) 
 * returns formatted start time of the given process */
PyObject *
getstart(proc_t *buf)
{
   time_t start;
   unsigned long seconds_since_boot = uptime(0,0);
   unsigned long seconds_since_1970 = time(NULL);
   unsigned long time_of_boot = seconds_since_1970 - seconds_since_boot;
   struct tm *tm;
   
   start = time_of_boot + buf->start_time / Hertz;

   buf->start_time = start;

   tm = gmtime(&start);

   if(tm == NULL)
      return NULL;

   return Py_BuildValue("(iiiiiiiii)",
                                tm->tm_year+1900,
                                tm->tm_mon+1,
                                tm->tm_mday,
                                tm->tm_hour,
                                tm->tm_min,
                                tm->tm_sec,
                                tm->tm_wday,
                                tm->tm_yday+1,
                                tm->tm_isdst);
}

/* ps's TIME parameter */
PyObject *
gettime(proc_t *buf)
{
   unsigned long long t;
   unsigned u;

   t = buf->utime + buf->stime;
   u = t / Hertz;

   return Py_BuildValue("(II)", u/60U, u%60U);
}

PyObject *
fill_buf(proc_t *buf, char *tty, PyObject *cmdline)
{
   PyObject *o;
   PyObject *pcpustr;
   
   pcpustr = setpcpu(buf);

   // when changing something in this dictionary, please remember to make
   // appropriate changes in _procpy.h and procpy.py!
   o = Py_BuildValue("{"
                        "s:i,s:i,"
                        "s:I,"
                        "s:c,"
                        "s:K,s:K,s:K,s:K,s:K,"
                        "s:l,s:l,s:l"
                        "s:s,s:s,s:s,s:s,s:s,"
                        "s:s,s:s,s:s,s:s,"
                        "s:i,s:i,s:i,s:i,s:i,"
                        "s:i,s:i,s:i,s:i,s:i,"
                        "s:i,s:i,s:i,"
                        "s:O,"

                        // procpy specific
                        "s:O,s:O,"
                        "s:s,"
                        "s:O,s:O,s:O"
                      "}",
#if defined(PROCPSVER30) || defined(PROCPSVER31)
                       "tid", buf->pid,         // s:i
#else
                       "tid", buf->tid,         // s:i
#endif
                      "ppid", buf->ppid,        // s:i
                      "pcpu", buf->pcpu,        // s:I
                     "state", buf->state,       // s:c
                     "utime", buf->utime,       // s:K
                     "stime", buf->stime,       // s:K
                    "cutime", buf->cutime,      // s:K
                    "cstime", buf->cstime,      // s:K
                "start_time", buf->start_time,  // s:K
                  "priority", buf->priority,    // s:l
                      "nice", buf->nice,        // s:l
                       "rss", buf->rss,         // s:l
                     "euser", buf->euser,       // s:s
                     "ruser", buf->ruser,       // s:s
                     "suser", buf->suser,       // s:s
                     "fuser", buf->fuser,       // s:s
                    "egroup", buf->egroup,      // s:s
                    "rgroup", buf->rgroup,      // s:s
                    "sgroup", buf->sgroup,      // s:s
                    "fgroup", buf->fgroup,      // s:s
                       "cmd", buf->cmd,         // s:s
                      "pgrp", buf->pgrp,        // s:i
                   "session", buf->session,     // s:i
#if defined(PROCPSVER30) || defined(PROCPSVER31)
                   "nlwp", 0,
#else

                      "nlwp", buf->nlwp,        // s:i
#endif

#ifndef PROCPSVER30
                      "tgid", buf->tgid,        // s:i
#else
                      "tgid", 0,
#endif
                       "tty", buf->tty,         // s:i
                      "euid", buf->euid,        // s:i
                      "egid", buf->egid,        // s:i
                      "ruid", buf->ruid,        // s:i
                      "rgid", buf->rgid,        // s:i
                      "suid", buf->sgid,        // s:i
                     "tpgid", buf->tpgid,       // s:i
               "exit_signal", buf->exit_signal, // s:i
                 "processor", buf->processor,   // s:i
                   "environ", getenviron(buf),  // s:O

                 /* procpy specific */
                   "pcpustr", pcpustr,          // s:O
                   "pmemstr", getpmem(buf),     // s:O
                    "ttynam", tty,              // s:s
                   "cmdline", cmdline,          // s:O
                     "start", getstart(buf),    // s:O
                      "time", gettime(buf));    // s:O
   return o;
}

/* 
 * readproc_general()
 *
 * list_or_dict == 0 => return list
 * list_or_dict != 0 => return dictionary
 */
PyObject *
readproc_general(int list_or_dict)
{
   PyObject *o;
   PyObject *k, *v;
   PyObject *cmdline;
   PROCTAB *ptp;
   proc_t buf;
   char tetey[128] = "";
   int flags = setflags();
   int tid;

   if(list_or_dict == 0) {
      if((o = PyList_New(0)) == NULL)
         return NULL;
   }
   else {
      if((o = PyDict_New()) == NULL)
         return NULL;
   }

   if(!(ptp = openproc(flags))) {
      PyErr_SetString(ProcpyError, "Cannot access /proc.\n");
      return NULL;
   }

   while(readproc(ptp, &buf)) {
      memset(&tetey, 0, sizeof(tetey));

#if defined(PROCPSVER30) || defined(PROCPSVER31)
      tid = buf.pid;
#else
      tid = buf.tid;
#endif

      if(buf.cmdline)
         cmdline = fillcmdline(buf.cmdline, tid);
      else
         continue;
 
      if(!tid)
         tid=-1;

      dev_to_tty(tetey, sizeof(tetey), buf.tty, tid, ABBREV_DEV);
      v = fill_buf(&buf, tetey, cmdline);
      if(v == NULL)
         return NULL;

      if(list_or_dict == 0) {
         if(PyList_Append(o, v))
            return NULL;
         Py_DECREF(v);
      }
      else {
         k = PyInt_FromLong((long)tid);
         if(PyDict_SetItem(o, k, v) < 0) {
            Py_DECREF(k);
            Py_DECREF(v);
            return NULL;
         }
         Py_DECREF(k);
         Py_DECREF(v);
      }
   }
   //freeproc(buf);
   closeproc(ptp);
   return o;
}


/*
 * readproc()
 */

PyDoc_STRVAR(readproc_doc,
   "readproc() -> returns list of dictionaries\n\n"
   PROCT_DOC);

PyObject *
_procpy_readproc(PyObject *self, PyObject *args)
{
   if(!PyArg_ParseTuple(args, ""))
      return NULL;

   return readproc_general(0);
}


/*
 * readproc_dict()
 */

PyDoc_STRVAR(readproc_dict_doc,
   "readproc() -> returns dictionary of dictionaries\n\n"
   PROCT_DOC);

PyObject *
_procpy_readproc_dict(PyObject *self, PyObject *args)
{
   if(!PyArg_ParseTuple(args, ""))
     return NULL;

   return readproc_general(1);
}


/*
 * readproc_by_pid()
 */

PyDoc_STRVAR(readproc_by_pid,
   "readproc_by_pid(pid) -> returns a dictionary\n\n"
   PROCT_DOC);

PyObject *
_procpy_readproc_by_pid(PyObject *self, PyObject *args)
{  
   PyObject *o;
   PyObject *cmdline;
   PROCTAB *ptp;
   proc_t buf;
   int i[1];
   char tetey[128];
   int flags = setflags();
   int tid;

   flags |= PROC_PID;

   if(!PyArg_ParseTuple(args, "i", &i[0]))
     return NULL;

   memset(&buf, 0, sizeof(buf));

   if(!(ptp = openproc(flags, i))) {
      PyErr_SetString(ProcpyError, "Cannot access /proc.\n");
      return NULL;
   }

   if(readproc(ptp, &buf) == NULL) {
      PyErr_SetString(ProcpyError, "PID not found.\n");
      return NULL;
   }

#if defined(PROCPSVER30) || defined(PROCPSVER31)
   tid = buf.pid;
#else
   tid = buf.tid;
#endif

   cmdline = fillcmdline(buf.cmdline, tid);

   if(!tid) {
      PyErr_SetString(ProcpyError, "PID not found.\n");
      return NULL;
   }
   dev_to_tty(tetey, sizeof(tetey), buf.tty, tid, ABBREV_DEV);
   o = fill_buf(&buf, tetey, cmdline);

   if(o == NULL)
      return NULL;

   //freeproc(buf);   
   closeproc(ptp);  
   return o;
}


PyDoc_STRVAR(uptime_doc, "uptime() -> uptime");

PyObject *
_procpy_uptime(PyObject *self, PyObject *args)
{
   if(!PyArg_ParseTuple(args, ""))
      return NULL;

   return Py_BuildValue("s", sprint_uptime());
}

PyDoc_STRVAR(loadavg_doc, "loadavg() -> returns a string;\n\n\
  Returned string contains system load averages for the past 1, 5,\n\
  and 15 minutes");

PyObject *
_procpy_loadavg(PyObject *self, PyObject *args)
{
   static double av[3];

   if(!PyArg_ParseTuple(args, ""))
      return NULL;

   loadavg(&av[0], &av[1], &av[2]);

   return Py_BuildValue("ddd", av[0], av[1], av[2]);
}

PyDoc_STRVAR(meminfo_doc, "meminfo() -> [MemTotal, MemUsed, MemFree, Buffers, SwapTotal, SwapUsed, SwapFree, Cached");

PyObject *
_procpy_meminfo(PyObject *self, PyObject *args)
{
   if(!PyArg_ParseTuple(args, ""))
      return NULL;

   meminfo();

   return Py_BuildValue("kkkkkkkk", kb_main_total, kb_main_used,
                                    kb_main_free, kb_main_buffers,
                                    kb_swap_total, kb_swap_used,
                                    kb_swap_free, kb_main_cached);
}


static PyMethodDef _procpyMethods[] = {
     {"loadavg", _procpy_loadavg, METH_VARARGS, loadavg_doc},
     {"meminfo", _procpy_meminfo, METH_VARARGS, meminfo_doc},
     {"readproc", _procpy_readproc, METH_VARARGS, readproc_doc},
     {"readproc_dict", _procpy_readproc_dict, METH_VARARGS, readproc_dict_doc},
     {"readproc_by_pid", _procpy_readproc_by_pid, METH_VARARGS,
                                                            readproc_by_pid},
     {"uptime", _procpy_uptime, METH_VARARGS, uptime_doc},
     {NULL, NULL, 0, NULL}
};


PyMODINIT_FUNC
init_procpy(void)
{
     PyObject *m, *d;

     m = Py_InitModule("_procpy", _procpyMethods);
     d = PyModule_GetDict(m);
     ProcpyError = PyErr_NewException("_procpy.error", NULL, NULL);
     Py_INCREF(ProcpyError);
     PyModule_AddObject(m, "error", ProcpyError);
}
